package dd.lo.pdfimg;

import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.beans.property.SimpleStringProperty;
import javafx.beans.property.StringProperty;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.control.TextField;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDDocumentInformation;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.PDResources;
import org.apache.pdfbox.pdmodel.common.PDStream;
import org.apache.pdfbox.pdmodel.graphics.PDXObject;
import org.apache.pdfbox.pdmodel.graphics.image.PDImageXObject;
import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.imgcodecs.Imgcodecs;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.*;

public class PdfImgExtractor extends Application {

    public static void start(String[] args) {
        launch(args);
    }

    private static final String OUTPUT_DIR = "/Users/kwok/Downloads/pdf_img/";

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    private static final int MAX_IMG_WIDTH = 108;
    private static final int MAX_IMG_HEIGHT = 77;

    private FlowPane imageFlowPane;

    private StringProperty gUrlProperty;

    private void pdfProcess(PDDocument document, String fileName) {
        int dotIndex = fileName.indexOf(".");
        if (dotIndex < 0) dotIndex = fileName.length();
        String prefix = fileName.substring(0, dotIndex).replaceAll("&", "").replaceAll("#", "").replaceAll("/", "").replaceAll(";", "");
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                int imageCnt = 0;
                for (int i = 0; i < document.getNumberOfPages(); ++i) {
                    System.out.printf("第%d页\n", i);
                    PDPage page = document.getPage(i);
                    PDResources resources = page.getResources();
                    for (COSName cosName : resources.getXObjectNames()) {
                        PDXObject xObject;
                        try {
                            xObject = resources.getXObject(cosName);
                        } catch (IOException e) {
                            continue;
                        }
                        if (xObject instanceof PDImageXObject) {
                            PDStream stream = xObject.getStream();
                            try {
                                //先转一次流才能拿到真正的图片对象
                                PDImageXObject imageObject = new PDImageXObject(stream, resources);
                                //获得awt的BufferedImage
                                BufferedImage image = imageObject.getImage();
                                //将BufferedImage利用ImageIO.write按jpeg格式输出到ByteArrayOutputStream流中
                                ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
                                ImageIO.write(image, "jpeg", byteArrayOutputStream);
                                byteArrayOutputStream.flush();
                                //此时就可以得到图片的byte信息，在ByteArrayOutputStream中，然后就可以通过byte信息将图片load到Mat中
                                Mat img = Imgcodecs.imdecode(new MatOfByte(byteArrayOutputStream.toByteArray()), Imgcodecs.IMREAD_UNCHANGED);
                                if (img.empty()) continue;
                                if (img.width() < 50 || img.height() < 50) continue; //太小的图片忽略掉
                                img.convertTo(img, CvType.CV_8UC3);
                                if (img.channels() < 3) {
                                    //如果图片不是3通道的，先自己合并自己变成3通道，否则后续步骤处理报错
                                    Core.merge(Arrays.asList(img, img, img), img);
                                }
                                //输出提取的图片
                                String outputNamePattern = "%s_%d_img_%d.jpeg";
                                String outputPath = String.format(outputNamePattern, OUTPUT_DIR + prefix, (imageCnt + 1), 0);
                                File outputFile = new File(outputPath);
                                if (!outputFile.exists()) {
                                    try {
                                        outputFile.createNewFile();
                                        Imgcodecs.imwrite(outputPath, img);
                                    } catch (IOException e) {
                                        e.printStackTrace();
                                    }
                                }
                                //寻找图片中的小图，若有，也输出文件
                                Mat[] cropResult = ImageCropper.cropImage(img);
                                for (int j = 0; j < cropResult.length; ++j) {
                                    outputPath = String.format(outputNamePattern, OUTPUT_DIR + prefix, (imageCnt + 1), j + 1);
                                    outputFile = new File(outputPath);
                                    if (!outputFile.exists()) {
                                        try {
                                            outputFile.createNewFile();
                                            Imgcodecs.imwrite(outputPath, cropResult[j]);
                                        } catch (IOException e) {
                                            e.printStackTrace();
                                        }
                                    }
                                }
                                if (img.width() > MAX_IMG_WIDTH) {
                                    ImageUtils.fitImgSize(img, MAX_IMG_WIDTH, MAX_IMG_HEIGHT);
                                }
                                MatOfByte matOfByte = new MatOfByte();
                                Imgcodecs.imencode(".jpeg", img, matOfByte);
                                ImageView imageView = new ImageView();
                                imageView.setImage(new Image(new ByteArrayInputStream(matOfByte.toArray())));
                                Platform.runLater(() -> imageFlowPane.getChildren().add(imageView));
                                ++imageCnt;
                            } catch (IOException ignored) {
                            }
                        }
                    }
                }
                try {
                    document.close();
                } catch (IOException ignored) {}
                System.out.printf("%d images are extracted\n", imageCnt);
                return null;
            }
        };
        new Thread(task).start();
    }

    @Override
    public void start(Stage stage) {
        imageFlowPane = new FlowPane();
        //File chooser
        Button imageChooser = new Button("选择PDF");
        imageChooser.setOnMouseClicked(e -> {
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("请选择PDF");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            fileChooser.getExtensionFilters().addAll(
                    new FileChooser.ExtensionFilter("PDF", "*.pdf")
            );
            File pdfFile = fileChooser.showOpenDialog(stage);
            if (null == pdfFile) {
                System.out.println("用户取消选中文件");
                return;
            }
            //清空显示在flowPane上面的图片
            Platform.runLater(() -> imageFlowPane.getChildren().clear());
            //读取一个PDF
            try {
                PDDocument document = PDDocument.load(pdfFile);
                pdfProcess(document, pdfFile.getName());
            } catch (IOException ex) {
                System.out.println("加载PDF文件出错，" + ex.getMessage());
                Platform.exit();
            }
        });
        //URL input box
        TextField urlInput = new TextField();
        gUrlProperty = new SimpleStringProperty();
        gUrlProperty.bind(urlInput.textProperty());
        //Fetch from URL button
        Button fetchBtn = new Button("从网路下载");
        fetchBtn.setOnMouseClicked(event -> {
            String urlPath = gUrlProperty.get();
            if (null == urlPath || urlPath.indexOf("http") != 0) {
                return;
            }
            //清空显示在flowPane上面的图片
            Platform.runLater(() -> imageFlowPane.getChildren().clear());
            try {
                URL url = new URL(urlPath);
                PDDocument document = PDDocument.load(url.openStream());
                String title;
                PDDocumentInformation documentInformation = document.getDocumentInformation();
                if (null != documentInformation && null != documentInformation.getTitle() && documentInformation.getTitle().length() > 0) {
                    title = documentInformation.getTitle();
                } else {
                    title = new Date().toString();
                }
                pdfProcess(document, title);
            } catch (Exception e) {
                e.printStackTrace();
            }
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(imageChooser, urlInput, fetchBtn);
        bottomControl.setPrefHeight(50);
        BorderPane root = new BorderPane();
        root.setCenter(imageFlowPane);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("PDF image extractor");
        stage.setResizable(false);
        stage.show();
    }
}
